# -*- coding: utf-8 -*-

import torch
import segmentation_models_pytorch as smp
from inception import InceptionV4
from CC_CCI_dataset import DiagnoseDataset


class Accuracy(torch.nn.Module):
    def __init__(self, threshold=0.5, activation=None, ignore_channels=None):
         super(Accuracy, self).__init__()
         self.accuracy = smp.utils.metrics.Accuracy(threshold, activation, ignore_channels)
         self.__name__ = "accuracy"
         
    def forward(self, y_pr, y_gt):
        y_gt = y_gt.to(y_pr.device)
        y_pr = torch.argmax(y_pr, dim=1)
        return self.accuracy.forward(y_pr, y_gt)

class CrossEntropyLoss(torch.nn.Module):
    def __init__(self, device='cuda'):
        super(CrossEntropyLoss, self).__init__()
        self.__name__ = "CrossEntropyLoss"
    
    def forward(self, y, y_pred):
        return torch.nn.CrossEntropyLoss()(y, y_pred)


model = InceptionV4()
model.to('cuda')

CSV_PATH = r'D:\data\metadata.csv'
SCAN_DIR = r'D:\data\NCP'

dataset = DiagnoseDataset(CSV_PATH, SCAN_DIR, device='cuda')
trainset, validset = torch.utils.data.random_split(dataset, [int(len(dataset)*0.7), len(dataset)-int(len(dataset)*0.7)])
train_loader = torch.utils.data.DataLoader(trainset, batch_size=5, shuffle=True, drop_last=True)
valid_loader = torch.utils.data.DataLoader(validset, batch_size=1, shuffle=True, drop_last=True)

loss = CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters())
metrics = [
    Accuracy(),
]

train_epoch = smp.utils.train.TrainEpoch(
    model,
    loss=loss,
    metrics=metrics,
    optimizer=optimizer,
    device='cuda',
    verbose=True,
)

valid_epoch = smp.utils.train.ValidEpoch(
    model,
    loss=loss,
    metrics=metrics,
    device='cuda',
    verbose=True,
)

for i in range(0, 10):
    print('\nEpoch: {}'.format(i))
    train_logs = train_epoch.run(train_loader)
    valid_logs = valid_epoch.run(valid_loader)

torch.save(model, "classifier.pkl")